#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
两段式去水印和图像恢复网络
包含：
1. Steganographic Attack Network (去水印网络)
2. Image Restoration Network (图像恢复网络)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class ConvBlock(nn.Module):
    """基础卷积块"""
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, use_bn=True, activation='relu'):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.use_bn = use_bn
        if use_bn:
            self.bn = nn.BatchNorm2d(out_channels)
        
        if activation == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif activation == 'leaky_relu':
            self.activation = nn.LeakyReLU(0.2, inplace=True)
        elif activation == 'tanh':
            self.activation = nn.Tanh()
        else:
            self.activation = nn.Identity()
    
    def forward(self, x):
        x = self.conv(x)
        if self.use_bn:
            x = self.bn(x)
        x = self.activation(x)
        return x


class ResidualBlock(nn.Module):
    """真正的残差块：实现 output = input + F(input)"""
    def __init__(self, channels, kernel_size=3, padding=1):
        super(ResidualBlock, self).__init__()
        # 第一个卷积层
        self.conv1 = ConvBlock(channels, channels, kernel_size, padding=padding)
        # 第二个卷积层（不使用激活函数）
        self.conv2 = ConvBlock(channels, channels, kernel_size, padding=padding, activation='none')
        # 最终的激活函数
        self.final_activation = nn.ReLU(inplace=True)
    
    def forward(self, x):
        # 保存输入作为残差连接
        residual = x
        
        # 通过两个卷积层
        out = self.conv1(x)
        out = self.conv2(out)
        
        # 残差连接：output = input + F(input)
        out = out + residual
        
        # 应用最终激活函数
        out = self.final_activation(out)
        
        return out


class SteganographicAttackNetwork(nn.Module):
    """
    隐写攻击网络（去水印网络）
    目标：模拟隐写信号并施加于水印图像，干扰原有水印
    """
    def __init__(self, input_channels=3, base_channels=64):
        super(SteganographicAttackNetwork, self).__init__()
        
        # 编码器：提取水印特征
        self.encoder = nn.Sequential(
            ConvBlock(input_channels, base_channels, kernel_size=7, padding=3),
            ConvBlock(base_channels, base_channels * 2, kernel_size=4, stride=2, padding=1),
            ConvBlock(base_channels * 2, base_channels * 4, kernel_size=4, stride=2, padding=1),
        )
        
        # 残差块：学习水印分布规律
        self.residual_blocks = nn.Sequential(
            ResidualBlock(base_channels * 4),
            ResidualBlock(base_channels * 4),
            ResidualBlock(base_channels * 4),
        )
        
        # 解码器：生成攻击信号
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(base_channels * 4, base_channels * 2, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(base_channels * 2),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(base_channels * 2, base_channels, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(base_channels),
            nn.ReLU(inplace=True),
            ConvBlock(base_channels, input_channels, kernel_size=7, padding=3, use_bn=False, activation='tanh')
        )
        
        # 攻击强度调节
        self.attack_strength = nn.Parameter(torch.tensor(0.1))
        
    def forward(self, watermarked_image):
        """
        前向传播
        Args:
            watermarked_image: 水印图像 (s)
        Returns:
            attacked_image: 攻击后的图像 (a)
            attack_noise: 攻击噪声
        """
        # 编码
        encoded = self.encoder(watermarked_image)
        
        # 残差学习
        residual = self.residual_blocks(encoded)
        
        # 解码生成攻击信号
        attack_noise = self.decoder(residual)
        
        # 应用攻击信号
        attacked_image = watermarked_image + self.attack_strength * attack_noise
        attacked_image = torch.clamp(attacked_image, -1, 1)
        
        return attacked_image, attack_noise


class ImageRestorationNetwork(nn.Module):
    """
    图像恢复网络
    目标：通过优化网络计算残差图像，恢复高质量图像
    """
    def __init__(self, input_channels=3, base_channels=64):
        super(ImageRestorationNetwork, self).__init__()
        
        # 特征提取层
        self.feature_extractor = nn.Sequential(
            ConvBlock(input_channels, base_channels, kernel_size=3, padding=1),
            ConvBlock(base_channels, base_channels, kernel_size=3, padding=1),
        )
        
        # 深度残差学习
        self.deep_residual = nn.Sequential(
            ResidualBlock(base_channels),
            ResidualBlock(base_channels),
            ResidualBlock(base_channels),
            ResidualBlock(base_channels),
        )
        
        # 残差预测网络
        self.residual_predictor = nn.Sequential(
            ConvBlock(base_channels, base_channels // 2, kernel_size=3, padding=1),
            ConvBlock(base_channels // 2, input_channels, kernel_size=3, padding=1, use_bn=False, activation='tanh')
        )
        
    def forward(self, attacked_image):
        """
        前向传播
        Args:
            attacked_image: 攻击后的图像 (a)
        Returns:
            restored_image: 恢复后的图像 (r)
            residual_image: 残差图像 (R)
        """
        # 特征提取
        features = self.feature_extractor(attacked_image)
        
        # 深度残差学习
        deep_features = self.deep_residual(features)
        
        # 残差预测
        residual_image = self.residual_predictor(deep_features)
        
        # 最终图像恢复
        restored_image = attacked_image + residual_image
        restored_image = torch.clamp(restored_image, -1, 1)
        
        return restored_image, residual_image


class TwoStageWatermarkRemoval(nn.Module):
    """
    两段式去水印网络
    结合隐写攻击网络和图像恢复网络
    """
    def __init__(self, input_channels=3, attack_base_channels=64, restore_base_channels=64):
        super(TwoStageWatermarkRemoval, self).__init__()
        
        # 第一阶段：隐写攻击网络
        self.attack_network = SteganographicAttackNetwork(
            input_channels=input_channels,
            base_channels=attack_base_channels
        )
        
        # 第二阶段：图像恢复网络
        self.restoration_network = ImageRestorationNetwork(
            input_channels=input_channels,
            base_channels=restore_base_channels
        )
        
    def forward(self, watermarked_image, cover_image=None):
        """
        前向传播
        Args:
            watermarked_image: 水印图像 (s)
            cover_image: 原始图像 (c) - 可选，用于计算损失
        Returns:
            results: 包含所有中间结果和最终结果的字典
        """
        # 第一阶段：隐写攻击
        attacked_image, attack_noise = self.attack_network(watermarked_image)
        
        # 第二阶段：图像恢复
        restored_image, residual_image = self.restoration_network(attacked_image)
        
        # 返回所有结果
        results = {
            'watermarked_image': watermarked_image,    # s
            'attacked_image': attacked_image,          # a
            'restored_image': restored_image,          # r
            'residual_image': residual_image,          # R
            'attack_noise': attack_noise,              # 攻击噪声
            'cover_image': cover_image                 # c
        }
        
        return results 